package com.ruoyi.wetest.manager;

import com.ruoyi.common.utils.ip.IpUtils;
import com.ruoyi.wetest.config.WeTestConfig;
import com.ruoyi.wetest.domain.*;
import com.ruoyi.wetest.mapper.WetestStressTransMetricMapper;
import com.ruoyi.wetest.service.IWetestProjectsService;
import com.ruoyi.wetest.service.IWetestStatLogService;
import com.ruoyi.wetest.service.IWetestStressRecordService;
import com.ruoyi.wetest.stats.MessageMetric;
import com.ruoyi.wetest.utils.IdGeneratorSnowflake;
import com.ruoyi.wetest.utils.WeTestRet;
import com.ruoyi.wetest.vo.Pair;
import com.ruoyi.wetest.vo.WeTestStressInfoVo;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Service;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * @author Allen
 * @description TODO
 * @Date 2024/4/26 20:02
 */
@Service
public class WetestManager {
    protected final Logger logger = LoggerFactory.getLogger(this.getClass());

    @Autowired
    private IWetestProjectsService wetestProjectsService;
    @Autowired
    private IWetestStressRecordService iWetestStressRecordService;
    @Autowired
    private WetestStressTransMetricMapper wetestStressTransMetricMapper;
    @Autowired
    private IWetestStatLogService iWetestStatLogService;


    @Autowired
    private WeTestConfig weTestConfig;

    @Autowired
    private ThreadPoolTaskExecutor executor;
    @Autowired
    private ScheduledExecutorService sexecutor;


    @PostConstruct
    public void init() {
        logger.debug("init");
        sexecutor.scheduleAtFixedRate(this::fixedS, 5, 5, TimeUnit.SECONDS);

    }

    private void fixedS() {
        for (WeTestStressInfoVo value : currentStressProjectMap.values()) {
            MessageMetric metric = value.getMessageMetric();
            if (metric == null || !metric.isRunning()) continue;
            metric.fixedTimeSnapshot();
        }
    }

    private final ConcurrentHashMap<String, WeTestStressInfoVo> currentStressProjectMap = new ConcurrentHashMap<>();
    private final ConcurrentHashMap<String, String> currentStressKey = new ConcurrentHashMap<>();

    public WeTestStressInfoVo getWeTestStressInfoVoNullable(String testId) {
        String pid = currentStressKey.get(testId);
        if (pid == null) return null;
        return currentStressProjectMap.get(pid);
    }

    public WeLoadMasterDataCollectStartResp start(WeLoadMasterDataCollectStartReq req) {
        WeTestStressInfoVo ws = currentStressProjectMap.get(req.getProjectid());
        if (ws == null) {
            ws = new WeTestStressInfoVo(req.getProjectid(), "" + IdGeneratorSnowflake.getInstance().snowflakeId(), "" + IdGeneratorSnowflake.getInstance().snowflakeId());
            currentStressProjectMap.put(req.getProjectid(), ws);
            currentStressKey.put(ws.getTestId(), ws.getProjectid());
            WetestProjects projects = wetestProjectsService.selectWetestProjectsByProjectId(req.getProjectid());
            if (projects == null) {
                projects = new WetestProjects();
                projects.setProjectId(req.getProjectid());
                projects.setCurrentClient(0);
                projects.setProdectName(req.getProjectid());
                projects.setStatus(0);
                projects.setStressTimes(0);
                wetestProjectsService.insertWetestProjects(projects);
            }
        }
        int instanceId = ws.getInstanceId().incrementAndGet();
        ws.incClientNum();
        String ip = IpUtils.getIpAddr();
        ws.getInstancesIp().add(ip);
        String acAddr = weTestConfig.getAcAddr(req.getProjectid());
        logger.warn("start projectId:{},testId:{},instanceId:{},acAddr:{},enabledRecordTrans:{}", req.getProjectid(), ws.getTestId(), instanceId, acAddr, weTestConfig.isEnabledRecordTrans());
        return new WeLoadMasterDataCollectStartResp(acAddr, instanceId, ws.getTestId(), ws.getToken());
    }

    public int registerLoad(WeTestGaspReq req) {
        WeTestStressInfoVo ws = getWeTestStressInfoVoNullable(req.getTestid());
        if (ws == null) return WeTestRet.WS_NOT_FOUND;
        WetestProjects projects = wetestProjectsService.selectWetestProjectsByProjectId(ws.getProjectid());
        if (projects != null) {
            projects.setStatus(1);
            projects.setCurrentClient(ws.getClientNum());

            wetestProjectsService.updateWetestProjects(projects);
            addOrUpdateRecord(ws);
        }
        return 0;
    }

    public void statisticsInfo(WetestStressStatistics req) {
        WeTestStressInfoVo ws = getWeTestStressInfoVoNullable(req.getTestId());
        if (ws == null) return;
        MessageMetric messageMetric = ws.getMessageMetric();
        if (messageMetric == null) {
            messageMetric = getMessageMetricAndCreate(ws);
        }
        messageMetric.statistics(req);
    }

    /**
     * 事务上报处理
     */
    public void transInfo(WeTestTransReq req, String testId) {
        WeTestStressInfoVo ws = getWeTestStressInfoVoNullable(testId);
        if (ws == null) return;
        MessageMetric messageMetric = ws.getMessageMetric();
        if (messageMetric == null) {
            messageMetric = getMessageMetricAndCreate(ws);
        }
        messageMetric.addReceivedMsgStat(req.getData());
    }

    private MessageMetric getMessageMetricAndCreate(WeTestStressInfoVo ws) {
        synchronized (ws) {
            MessageMetric messageMetric = ws.getMessageMetric();
            if (messageMetric == null) {
                messageMetric = new MessageMetric(ws.getTestId(), executor, iWetestStatLogService);
                ws.setMessageMetric(messageMetric);
            }
            return messageMetric;
        }
    }

    public void stop(String projectid) {
        WeTestStressInfoVo ws = currentStressProjectMap.get(projectid);
        if (ws == null) return;
        ws.decClientNum();
        int clientNum = ws.getClientNum();
        logger.warn("stop projectid：{},testId:{},clientNum:{}", projectid, ws.getTestId(), clientNum);
        if (clientNum <= 0) {
            currentStressKey.remove(ws.getTestId());
            currentStressProjectMap.remove(projectid);
            ws.stop();
            MessageMetric messageMetric = ws.getMessageMetric();
            if (messageMetric != null) {
                Map<String, WetestStressTransMetric> snapshotList = messageMetric.getSnapshotList();
                if (!snapshotList.isEmpty()) {
                    List<WetestStressTransMetric> lss = new ArrayList<>(snapshotList.values());
                    wetestStressTransMetricMapper.saveBatchWetestStressTransMetric(lss);
                }
                messageMetric.stop();
            }
            addOrUpdateRecord(ws);
        }

        WetestProjects projects = wetestProjectsService.selectWetestProjectsByProjectId(projectid);
        if (projects != null) {
            if (clientNum <= 0) {
                projects.setStressTimes(projects.getStressTimes() + 1);
                projects.setStatus(0);
            }
            projects.setCurrentClient(clientNum);
            wetestProjectsService.updateWetestProjects(projects);
        }

    }

    public WetestStressRecord addOrUpdateRecord(WeTestStressInfoVo ws) {
        WetestStressRecord wsr = iWetestStressRecordService.selectWetestStressRecordByTestId(ws.getTestId());
        if (wsr == null) {
            wsr = new WetestStressRecord();
            wsr.setTestId(ws.getTestId());
            wsr.setProjectId(ws.getProjectid());
            wsr.setStressClient(ws.getInstanceId().get());
            wsr.setStartTime(ws.getStartDate());
            wsr.setEndTime(ws.getStopDate());
            addOrUpdateRecordExtracted(ws, wsr);
            iWetestStressRecordService.insertWetestStressRecord(wsr);
        } else {
            addOrUpdateRecordExtracted(ws, wsr);
            iWetestStressRecordService.updateWetestStressRecord(wsr);
        }
        return wsr;
    }

    private static void addOrUpdateRecordExtracted(WeTestStressInfoVo ws, WetestStressRecord wsr) {
        int tpcu = 0;
        int opcu = 0;
        int allTps = 0;
        int transNum = 0;
        MessageMetric messageMetric = ws.getMessageMetric();
        if (messageMetric != null) {
            transNum = messageMetric.getCodeStatsMapSize() - 1;
            allTps = (int) messageMetric.allCodeStat.getRxtx().getMeanRate();
            Pair<Integer, Integer> pcu = messageMetric.getPcu();
            tpcu = pcu.first;
            opcu = pcu.second;
        }
        wsr.setTotalPcu(tpcu);
        wsr.setOnlinePcu(opcu);
        wsr.setAllTps(allTps);
        wsr.setTransNum(transNum);
        wsr.setEndTime(ws.getStopDate());
        if (ws.getStopDate()!=null){
            wsr.setRemark(Arrays.toString(ws.getInstancesIp().toArray()));
        }
    }

    public List<String> getStressList() {
        ArrayList<String> ls = new ArrayList<>();
        for (WeTestStressInfoVo infoVo : currentStressProjectMap.values()) {
            ls.add(infoVo.toString());
        }
        return ls;
    }

    public Collection<WetestStressTransMetric> getMetricSnapshot(String testId) {
        WeTestStressInfoVo ws = getWeTestStressInfoVoNullable(testId);
        if (ws == null) return new ArrayList<>();
        MessageMetric messageMetric = ws.getMessageMetric();
        if (messageMetric == null) return new ArrayList<>();
        return messageMetric.getSnapshotList().values();
    }

    @PreDestroy
    public void shutdownGracefully() {
        logger.error("shutdown  currentStress:{}", currentStressProjectMap.size());
        for (WeTestStressInfoVo infoVo : currentStressProjectMap.values()) {
            infoVo.resetClientNum();
            stop(infoVo.getProjectid());
        }
    }


}
